import requests
from typing import Dict, List, Union
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from .base import BaseModel

class DeepSeekModel(BaseModel):
    """DeepSeek模型实现，支持深度思考"""
    
    def __init__(self, api_key: str, model_name: str, api_url: str):
        super().__init__(api_key, model_name)
        self.api_url = api_url
        self.supports_thinking = True
        
        # 配置重试策略
        retry_strategy = Retry(
            total=3,  # 最大重试次数
            backoff_factor=1,  # 重试间隔
            status_forcelist=[429, 500, 502, 503, 504]  # 需要重试的HTTP状态码
        )
        
        # 创建会话并应用重试策略
        self.session = requests.Session()
        self.session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
    
    def generate_response(self, messages: List[Dict[str, str]], enable_thinking: bool = False) -> Union[str, Dict[str, str]]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [{"role": msg["role"], "content": msg["content"]} for msg in messages],
            "max_tokens": 2000,
            "temperature": 0.7,
            "stream": False
        }
        
        if enable_thinking:
            data["reasoning_effort"] = "medium"
        
        try:
            response = self.session.post(self.api_url, headers=headers, json=data, timeout=60)
            response.raise_for_status()
            
            result = response.json()
            content = result.get('choices', [{}])[0].get('message', {}).get('content', '').strip()
            
            if enable_thinking and result.get('choices', [{}])[0].get('message', {}).get('reasoning_content'):
                thinking = result.get('choices', [{}])[0].get('message', {}).get('reasoning_content', '')
                return {"thinking": thinking, "response": content}
            
            return content
            
        except requests.exceptions.Timeout:
            raise Exception("连接超时，请检查网络状态或稍后重试")
        except requests.exceptions.ConnectionError:
            raise Exception("连接失败，请检查网络连接")
        except requests.exceptions.HTTPError as e:
            status_code = e.response.status_code if e.response is not None else None
            if status_code == 401:
                raise Exception("API密钥无效或已过期")
            elif status_code == 429:
                raise Exception("请求过于频繁，请稍后重试")
            elif status_code and status_code >= 500:
                raise Exception("服务器暂时不可用，请稍后重试")
            else:
                raise Exception(f"HTTP请求失败: {str(e)}")
        except requests.exceptions.RequestException as e:
            raise Exception(f"请求失败: {str(e)}")
        except Exception as e:
            raise Exception(f"生成回答时出错: {str(e)}")